// sm4_demo.c
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <openssl/evp.h>
#include <openssl/rand.h>
#include <openssl/err.h>

#define SM4_KEY_SIZE 16  // 128 bits
#define SM4_BLOCK_SIZE 16

// 打印十六进制
void print_hex(const char *label, const unsigned char *data, size_t len) {
    printf("%s: ", label);
    for (size_t i = 0; i < len; i++) {
        printf("%02x", data[i]);
    }
    printf("\n");
}

// PKCS#7 填充（用于 CBC 模式）
void pkcs7_pad(const unsigned char *in, size_t in_len, unsigned char *out, size_t *out_len) {
    size_t pad_len = SM4_BLOCK_SIZE - (in_len % SM4_BLOCK_SIZE);
    memcpy(out, in, in_len);
    for (size_t i = 0; i < pad_len; i++) {
        out[in_len + i] = (unsigned char)pad_len;
    }
    *out_len = in_len + pad_len;
}

// PKCS#7 去填充
size_t pkcs7_unpad(unsigned char *data, size_t data_len) {
    if (data_len == 0 || data_len % SM4_BLOCK_SIZE != 0) return data_len;
    unsigned char pad = data[data_len - 1];
    if (pad < 1 || pad > SM4_BLOCK_SIZE) return data_len;
    for (size_t i = 0; i < pad; i++) {
        if (data[data_len - 1 - i] != pad) return data_len;
    }
    return data_len - pad;
}

// SM4-ECB 加密
int sm4_ecb_encrypt(const unsigned char *key,
                    const unsigned char *plaintext, size_t pt_len,
                    unsigned char **ciphertext, size_t *ct_len) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    if (!ctx) return -1;

    if (EVP_EncryptInit_ex(ctx, EVP_sm4_ecb(), NULL, key, NULL) != 1) goto err;
    EVP_CIPHER_CTX_set_padding(ctx, 0); // ECB 模式通常不推荐填充，但为兼容可选

    // 由于 ECB 要求输入为块大小整数倍，我们强制填充（仅演示）
    size_t padded_len = ((pt_len + SM4_BLOCK_SIZE - 1) / SM4_BLOCK_SIZE) * SM4_BLOCK_SIZE;
    unsigned char *padded = calloc(1, padded_len);
    memcpy(padded, plaintext, pt_len);
    // 简单零填充（实际应使用 PKCS#7，但 ECB 本身不安全，仅演示）
    // 更推荐使用 CBC

    *ciphertext = calloc(1, padded_len + SM4_BLOCK_SIZE);
    int len1, len2;
    if (EVP_EncryptUpdate(ctx, *ciphertext, &len1, padded, padded_len) != 1) {
        free(padded);
        goto err;
    }
    if (EVP_EncryptFinal_ex(ctx, *ciphertext + len1, &len2) != 1) {
        free(padded);
        goto err;
    }
    *ct_len = len1 + len2;
    free(padded);
    EVP_CIPHER_CTX_free(ctx);
    return 0;

err:
    ERR_print_errors_fp(stderr);
    EVP_CIPHER_CTX_free(ctx);
    return -1;
}

// SM4-CBC 加密（推荐）
int sm4_cbc_encrypt(const unsigned char *key,
                    const unsigned char *plaintext, size_t pt_len,
                    unsigned char **ciphertext, size_t *ct_len,
                    unsigned char iv[SM4_BLOCK_SIZE]) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    if (!ctx) return -1;

    if (EVP_EncryptInit_ex(ctx, EVP_sm4_cbc(), NULL, key, iv) != 1) goto err;
    EVP_CIPHER_CTX_set_padding(ctx, 1); // 启用 PKCS#7 填充

    size_t buf_len = pt_len + SM4_BLOCK_SIZE;
    *ciphertext = calloc(1, buf_len);
    int len1, len2;

    if (EVP_EncryptUpdate(ctx, *ciphertext, &len1, plaintext, pt_len) != 1) goto err;
    if (EVP_EncryptFinal_ex(ctx, *ciphertext + len1, &len2) != 1) goto err;

    *ct_len = len1 + len2;
    EVP_CIPHER_CTX_free(ctx);
    return 0;

err:
    ERR_print_errors_fp(stderr);
    EVP_CIPHER_CTX_free(ctx);
    return -1;
}

// SM4-CBC 解密
int sm4_cbc_decrypt(const unsigned char *key,
                    const unsigned char *ciphertext, size_t ct_len,
                    unsigned char **plaintext, size_t *pt_len,
                    unsigned char iv[SM4_BLOCK_SIZE]) {
    EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
    if (!ctx) return -1;

    if (EVP_DecryptInit_ex(ctx, EVP_sm4_cbc(), NULL, key, iv) != 1) goto err;
    EVP_CIPHER_CTX_set_padding(ctx, 1);

    *plaintext = calloc(1, ct_len);
    int len1, len2;
    if (EVP_DecryptUpdate(ctx, *plaintext, &len1, ciphertext, ct_len) != 1) goto err;
    if (EVP_DecryptFinal_ex(ctx, *plaintext + len1, &len2) != 1) goto err;

    *pt_len = pkcs7_unpad(*plaintext, len1 + len2);
    EVP_CIPHER_CTX_free(ctx);
    return 0;

err:
    ERR_print_errors_fp(stderr);
    EVP_CIPHER_CTX_free(ctx);
    return -1;
}

int main(void) {
    const unsigned char key[SM4_KEY_SIZE] = "1234567890abcdef"; // 16字节密钥
    const unsigned char plaintext[] = "This is a secret message for SM4!";
    size_t pt_len = sizeof(plaintext) - 1;

    printf("Original plaintext: %s\n", plaintext);
    print_hex("Key", key, SM4_KEY_SIZE);

    // === SM4-CBC 模式（推荐）===
    unsigned char iv[SM4_BLOCK_SIZE];
    if (!RAND_bytes(iv, SM4_BLOCK_SIZE)) {
        fprintf(stderr, "Failed to generate IV\n");
        return 1;
    }
    print_hex("IV (random)", iv, SM4_BLOCK_SIZE);

    unsigned char *ciphertext = NULL;
    size_t ct_len = 0;
    if (sm4_cbc_encrypt(key, plaintext, pt_len, &ciphertext, &ct_len, iv) != 0) {
        fprintf(stderr, "SM4-CBC encrypt failed\n");
        return 1;
    }
    print_hex("Ciphertext (CBC)", ciphertext, ct_len);

    unsigned char *decrypted = NULL;
    size_t dec_len = 0;
    if (sm4_cbc_decrypt(key, ciphertext, ct_len, &decrypted, &dec_len, iv) != 0) {
        fprintf(stderr, "SM4-CBC decrypt failed\n");
        free(ciphertext);
        return 1;
    }

    printf("Decrypted text: %.*s\n", (int)dec_len, decrypted);
    if (dec_len == pt_len && memcmp(plaintext, decrypted, pt_len) == 0) {
        printf("[SUCCESS] CBC: Decryption matches original!\n");
    } else {
        printf("[ERROR] CBC: Decryption mismatch!\n");
    }

    // === SM4-ECB 模式（仅演示，不推荐）===
    /*
    unsigned char *ecb_ct = NULL;
    size_t ecb_ct_len = 0;
    if (sm4_ecb_encrypt(key, plaintext, pt_len, &ecb_ct, &ecb_ct_len) == 0) {
        print_hex("Ciphertext (ECB)", ecb_ct, ecb_ct_len);
        free(ecb_ct);
    }
    */

    // 清理
    free(ciphertext);
    free(decrypted);

    return 0;
}
